"""
Author: Morphlng
Date: 2024-03-20 09:08:19
LastEditTime: 2024-03-20 14:29:53
LastEditors: Morphlng
Description: Test for Quaternion
FilePath: /DrivingGym/tests/common/geom/quaternion_test.py
"""

import math
import random

import numpy as np

from driving_gym.simulation.common.geom import Quaternion, Rotation, Transform


def is_close(a, b):
    return math.isclose(a + 1, b + 1)


class TestQuaternion:
    def test_initialization_normalization(self):
        q = Quaternion(1, 1, 1, 1)
        norm = math.sqrt(q.w**2 + q.x**2 + q.y**2 + q.z**2)
        assert is_close(norm, 1), "Quaternion is not normalized"

    def test_zero_quaternion(self):
        q = Quaternion(0, 0, 0, 0)
        assert (
            q.w == 0 and q.x == 0 and q.y == 0 and q.z == 0
        ), "Zero vector does not create a zero quaternion"

    def test_create_matrix(self):
        try:
            from scipy.spatial.transform import Rotation as R
        except ImportError:
            # Simple test only
            q = Quaternion(1, 0, 0, 0)
            trans = Transform(None, Rotation(0, 0, 0))
            expected_matrix = np.eye(3)
            np.testing.assert_allclose(
                q.matrix,
                expected_matrix,
                rtol=1e-5,
                err_msg="Matrix creation failed for no rotation",
            )
            np.testing.assert_allclose(
                q.transform_matrix,
                trans.get_matrix()[:3, :3],
                rtol=1e-5,
                err_msg="Transform matrix creation failed for no rotation",
            )
            return

        test_angles = [
            (
                random.uniform(-180, 180),
                random.uniform(-180, 180),
                random.uniform(-180, 180),
            )
            for _ in range(10)
        ]

        for pitch, yaw, roll in test_angles:
            rot = Rotation(pitch, yaw, roll)
            quat = Quaternion.from_rotation(rot)
            mat = quat.matrix
            scipy_rot = R.from_euler("xyz", [roll, pitch, yaw], degrees=True)
            scipy_mat = scipy_rot.as_matrix()

            np.testing.assert_allclose(
                mat,
                scipy_mat,
                rtol=1e-5,
                err_msg=f"Matrix creation failed for Euler angles {(pitch, yaw, roll)}",
            )
            np.testing.assert_allclose(
                quat.transform_matrix,
                Transform(None, rot).get_matrix()[:3, :3],
                rtol=1e-5,
                err_msg=f"Transform matrix creation failed for Euler angles {(pitch, yaw, roll)}",
            )

    def test_from_rotation(self):
        try:
            from scipy.spatial.transform import Rotation as R
        except ImportError:
            # Simple test only
            quat = Quaternion.from_rotation(Rotation(20, 30, 40))
            expect_quat = np.array([0.07043934, 0.2968829, 0.28311405, 0.90925534])
            np.testing.assert_allclose(
                [quat.x, quat.y, quat.z, quat.w],
                expect_quat,
                rtol=1e-5,
                err_msg="Quaternion conversion failed for Euler angles (20, 30, 40)",
            )
            return

        test_angles = [
            (
                random.uniform(-180, 180),
                random.uniform(-180, 180),
                random.uniform(-180, 180),
            )
            for _ in range(10)
        ]

        for pitch, yaw, roll in test_angles:
            # Create your Quaternion from Euler angles
            quat = Quaternion.from_rotation(Rotation(pitch, yaw, roll))

            # Create scipy Rotation from Euler angles and convert to quaternion
            scipy_rot = R.from_euler("xyz", [roll, pitch, yaw], degrees=True)
            # scipy's quaternion is [x, y, z, w]
            scipy_quat = scipy_rot.as_quat()

            # Compare the quaternions
            np.testing.assert_allclose(
                [
                    quat.x,
                    quat.y,
                    quat.z,
                    quat.w,
                ],
                scipy_quat,
                rtol=1e-5,
                err_msg=f"Quaternion conversion failed for Euler angles {(pitch, yaw, roll)}",
            )

            # Convert back to Euler angles and compare
            euler = quat.as_rotation()
            scipy_euler = scipy_rot.as_euler("xyz", degrees=True)

            # Compare the Euler angles
            np.testing.assert_allclose(
                [euler.roll, euler.pitch, euler.yaw],
                scipy_euler,
                rtol=1e-5,
                err_msg=f"Euler angles conversion failed for Quaternion {quat}",
            )

    def test_multiplication(self):
        q1 = Quaternion(1, 0, 0, 0)
        q2 = Quaternion(0, 1, 0, 0)
        q_product = q1 * q2

        expected_q = Quaternion(0, 1, 0, 0)
        assert is_close(q_product.w, expected_q.w) and is_close(
            q_product.x, expected_q.x
        ), "Multiplication failed"
